// Copyright (c) 2021 Presto Labs Pte. Ltd.
// Author: chunhui

#pragma once

#include <cstdint>
#include <functional>
#include <unordered_map>
#include <utility>
#include <vector>

#include <fmt/format.h>
#include <glog/logging.h>

namespace coin2::strategy::util {

class TimeIndexedHeap {
 public:
  explicit TimeIndexedHeap(const std::function<bool(double, double)>& comp) : comp_(comp) {}

  static bool Less(double x, double y) { return x < y; }
  static bool Greater(double x, double y) { return x > y; }

  void Pop() {
    CHECK_GT(Size(), 0);
    Remove(id_to_ts_[0]);
  }

  void Add(int64_t ts, double val) {
    if (std::isfinite(val)) {
      ts_to_id_[ts] = Size();
      id_to_ts_.push_back(ts);
      array_.push_back(val);
      ShiftUp(array_.size() - 1);
    }
  }

  std::pair<int64_t, double> Top() {
    CHECK_GT(Size(), 0);
    return {id_to_ts_.front(), array_.front()};
  }

  size_t Size() const { return array_.size(); }
  bool Empty() const { return Size() == 0; }
  inline bool Contains(int64_t ts) const { return ts_to_id_.find(ts) != ts_to_id_.end(); }

  void Update(int64_t ts, double value) {
    if (std::isfinite(value) && Contains(ts)) {
      size_t i = ts_to_id_[ts];
      CHECK(i < array_.size());
      array_[i] = value;
      ShiftUp(i);
      ShiftDown(i);
    } else if (std::isfinite(value)) {
      Add(ts, value);
    }
  }

  void Remove(int64_t ts) {
    if (Contains(ts)) {
      size_t i = ts_to_id_[ts];
      CHECK(i < array_.size());
      Swap(i, array_.size() - 1);
      array_.pop_back();
      int64_t ts_to_remove = id_to_ts_.back();
      ts_to_id_.erase(ts_to_remove);
      id_to_ts_.pop_back();
      ShiftUp(i);
      ShiftDown(i);
    }
  }

 private:
  std::vector<double> array_;
  std::vector<int64_t> id_to_ts_;
  std::unordered_map<int64_t, size_t> ts_to_id_;
  std::function<bool(double, double)> comp_;

  inline static size_t GetLeftChild(size_t i) { return i * 2 + 1; }
  inline static size_t GetRightChild(size_t i) { return i * 2 + 2; }
  inline static size_t GetParent(size_t i) {
    if (i > 0) {
      return (i - 1) / 2;
    }
    return 0;
  }
  inline size_t ChooseChild(size_t i) const {
    size_t left = GetLeftChild(i);
    size_t right = GetRightChild(i);
    if (left < array_.size() && right < array_.size()) {
      if (comp_(array_[left], array_[right])) {
        return left;
      } else {
        return right;
      }
    } else {
      return left;
    }
  }

  inline void Swap(size_t i, size_t j) {
    double tmp = array_[i];
    array_[i] = array_[j];
    array_[j] = tmp;
    int64_t ts_i = id_to_ts_[i];
    int64_t ts_j = id_to_ts_[j];
    ts_to_id_[ts_i] = j;
    ts_to_id_[ts_j] = i;
    id_to_ts_[i] = ts_j;
    id_to_ts_[j] = ts_i;
  }

  void ShiftUp(size_t i) {
    size_t p = GetParent(i);
    while (i > 0 && comp_(array_[i], array_[p])) {
      Swap(p, i);
      i = p;
      p = GetParent(i);
    }
  }

  void ShiftDown(size_t i) {
    size_t child = ChooseChild(i);
    while (child < array_.size() && comp_(array_[child], array_[i])) {
      Swap(child, i);
      i = child;
      child = ChooseChild(i);
    }
  }
};

}  // namespace coin2::strategy::util
